"""SSH Slurm MCP Server - Main Server Implementation"""

import asyncio
import logging
import os
import sys
from typing import Dict, Any, Optional
import argparse
from pathlib import Path

from fastmcp import FastMCP
from dotenv import load_dotenv
from src.ssh_slurm_tool.ssh_manager import SSHManager
from src.ssh_slurm_tool.slurm_manager import SlurmManager
from src.ssh_slurm_tool.compiler import CompilerManager

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

# Load environment variables from .env file
env_path = Path(__file__).parent / '.env'
if env_path.exists():
    load_dotenv(env_path)
    logger.info(f"Loaded environment variables from {env_path}")
else:
    logger.warning(f"No .env file found at {env_path}")

# Create MCP server instance
mcp = FastMCP("SSH Slurm Programming Server")

# Initialize managers (will be done when server starts)
ssh_manager: Optional[SSHManager] = None
slurm_manager: Optional[SlurmManager] = None
compiler_manager: Optional[CompilerManager] = None

# Default SSH connection settings from environment variables
DEFAULT_SSH_HOST = os.getenv('SSH_HOST', '')
DEFAULT_SSH_USERNAME = os.getenv('SSH_USERNAME', '')
DEFAULT_SSH_KEY_PATH = os.getenv('SSH_KEY_PATH', '')
DEFAULT_SSH_PASSWORD = os.getenv('SSH_PASSWORD', '')
DEFAULT_SSH_PORT = int(os.getenv('SSH_PORT', '22'))


def initialize_managers():
    """Initialize all manager instances"""
    global ssh_manager, slurm_manager, compiler_manager
    
    ssh_manager = SSHManager()
    slurm_manager = SlurmManager(ssh_manager)
    compiler_manager = CompilerManager(ssh_manager)
    
    logger.info("Initialized SSH Slurm MCP Server managers")


def get_or_create_connection(connection_id: str = "") -> str:
    """Get existing connection or auto-create one using .env settings
    
    Args:
        connection_id: Specific connection ID to get, or empty to auto-connect
        
    Returns:
        Valid connection ID
        
    Raises:
        Exception: If no connection exists and auto-connect fails
    """
    if not ssh_manager:
        raise Exception("SSH manager not initialized")
    
    # If connection_id provided, try to use it
    if connection_id:
        if connection_id in ssh_manager.clients:
            return connection_id
        else:
            raise Exception(f"No active connection found: {connection_id}")
    
    # No connection_id provided - check for existing connections
    if ssh_manager.clients:
        # Return the first available connection
        return list(ssh_manager.clients.keys())[0]
    
    # No connections exist - try auto-connect using .env settings
    if not DEFAULT_SSH_HOST or not DEFAULT_SSH_USERNAME:
        missing = []
        if not DEFAULT_SSH_HOST: missing.append("SSH_HOST")
        if not DEFAULT_SSH_USERNAME: missing.append("SSH_USERNAME")
        raise Exception(f"No active connections and cannot auto-connect: Missing {', '.join(missing)} in .env file")
    
    if not DEFAULT_SSH_KEY_PATH and not DEFAULT_SSH_PASSWORD:
        raise Exception("No active connections and cannot auto-connect: Neither SSH_KEY_PATH nor SSH_PASSWORD set in .env file")
    
    # Auto-connect using .env settings
    logger.info("No active connections found. Auto-connecting using .env settings...")
    
    # Expand tilde in key path
    actual_key_path = DEFAULT_SSH_KEY_PATH
    if actual_key_path and actual_key_path.startswith('~'):
        actual_key_path = os.path.expanduser(actual_key_path)
    
    connection_id = ssh_manager.connect(
        host=DEFAULT_SSH_HOST,
        username=DEFAULT_SSH_USERNAME,
        key_path=actual_key_path if actual_key_path else None,
        password=DEFAULT_SSH_PASSWORD if DEFAULT_SSH_PASSWORD else None,
        port=DEFAULT_SSH_PORT
    )
    
    logger.info(f"Auto-connected to {connection_id}")
    return connection_id


@mcp.tool
def connect_ssh(host: str = "", username: str = "", key_path: str = "", 
                password: str = "", port: int = 0) -> str:
    """Establish SSH connection to remote server
    
    Automatically uses settings from .env file if parameters are not provided.
    
    Args:
        host: Remote server hostname or IP address (uses SSH_HOST from .env if empty)
        username: SSH username for authentication (uses SSH_USERNAME from .env if empty)
        key_path: Path to SSH private key file (uses SSH_KEY_PATH from .env if empty)
        password: SSH password (uses SSH_PASSWORD from .env if empty)
        port: SSH port number (uses SSH_PORT from .env if 0, defaults to 22)
        
    Returns:
        Connection ID for use in other operations
    """
    try:
        if not ssh_manager:
            raise Exception("SSH manager not initialized")
        
        # Use .env defaults if parameters are not provided
        actual_host = host or DEFAULT_SSH_HOST
        actual_username = username or DEFAULT_SSH_USERNAME
        actual_key_path = key_path or DEFAULT_SSH_KEY_PATH
        actual_password = password or DEFAULT_SSH_PASSWORD
        actual_port = port if port != 0 else DEFAULT_SSH_PORT
        
        # Expand tilde in key path
        if actual_key_path and actual_key_path.startswith('~'):
            actual_key_path = os.path.expanduser(actual_key_path)
        
        # Validate input
        if not actual_host or not actual_username:
            missing = []
            if not actual_host: missing.append("SSH_HOST")
            if not actual_username: missing.append("SSH_USERNAME") 
            raise ValueError(f"Missing required SSH connection info: {', '.join(missing)}. "
                           f"Please set these in .env file or provide as parameters.")
        
        if not actual_key_path and not actual_password:
            raise ValueError("Either SSH_KEY_PATH or SSH_PASSWORD must be set in .env file or provided as parameters")
        
        # Log connection attempt (without sensitive info)
        logger.info(f"Attempting SSH connection to {actual_username}@{actual_host}:{actual_port}")
        
        # Establish connection
        connection_id = ssh_manager.connect(
            host=actual_host,
            username=actual_username,
            key_path=actual_key_path if actual_key_path else None,
            password=actual_password if actual_password else None,
            port=actual_port
        )
        
        return f"Successfully connected to {connection_id}. Use this connection ID for subsequent operations."
        
    except Exception as e:
        logger.error(f"SSH connection failed: {str(e)}")
        return f"SSH connection failed: {str(e)}"


@mcp.tool
def auto_connect_ssh() -> str:
    """Auto-connect to SSH server using settings from .env file
    
    This is a convenience function that connects using all settings from .env file.
    Equivalent to calling connect_ssh() with no parameters.
    
    Returns:
        Connection ID for use in other operations, or error message
    """
    try:
        if not DEFAULT_SSH_HOST or not DEFAULT_SSH_USERNAME:
            missing = []
            if not DEFAULT_SSH_HOST: missing.append("SSH_HOST")
            if not DEFAULT_SSH_USERNAME: missing.append("SSH_USERNAME")
            return f"Cannot auto-connect: Missing {', '.join(missing)} in .env file. Please configure these settings."
        
        if not DEFAULT_SSH_KEY_PATH and not DEFAULT_SSH_PASSWORD:
            return "Cannot auto-connect: Neither SSH_KEY_PATH nor SSH_PASSWORD is set in .env file."
        
        # Call connect_ssh with no parameters to use all .env defaults
        result = connect_ssh()
        
        # Add information about which .env settings were used
        auth_method = "SSH key" if DEFAULT_SSH_KEY_PATH else "password"
        return f"""Auto-connected using .env settings:
• Host: {DEFAULT_SSH_USERNAME}@{DEFAULT_SSH_HOST}:{DEFAULT_SSH_PORT}
• Authentication: {auth_method}

{result}"""
        
    except Exception as e:
        logger.error(f"Auto-connect failed: {str(e)}")
        return f"Auto-connect failed: {str(e)}"


@mcp.tool
def disconnect_ssh(connection_id: str = "") -> str:
    """Disconnect from SSH server
    
    Args:
        connection_id: Specific connection to disconnect (empty = disconnect all)
        
    Returns:
        Disconnection status message
    """
    try:
        if not ssh_manager:
            raise Exception("SSH manager not initialized")
        
        if connection_id:
            success = ssh_manager.disconnect(connection_id)
            if success:
                return f"Successfully disconnected from {connection_id}"
            else:
                return f"Failed to disconnect from {connection_id} (connection may not exist)"
        else:
            ssh_manager.disconnect()
            return "Successfully disconnected from all SSH connections"
            
    except Exception as e:
        logger.error(f"SSH disconnection failed: {str(e)}")
        return f"SSH disconnection failed: {str(e)}"


@mcp.tool
def compile_code(connection_id: str = "", source_file: str = "", compiler: str = "",
                output_file: str = "", compile_flags: str = "", working_dir: str = "~") -> str:
    """Compile source code on remote server
    
    Args:
        connection_id: SSH connection identifier (empty = auto-connect using .env)
        source_file: Path to source file on remote server
        compiler: Compiler to use (gcc, hipcc, python)
        output_file: Output executable name (auto-generated if empty)
        compile_flags: Additional compiler flags (optional)
        working_dir: Working directory for compilation (default: ~)
        
    Returns:
        Compilation result and status
    """
    try:
        if not compiler_manager:
            raise Exception("Compiler manager not initialized")
        
        # Get or create connection
        actual_connection_id = get_or_create_connection(connection_id)
        
        result = compiler_manager.compile_code(
            connection_id=actual_connection_id,
            source_file=source_file,
            compiler=compiler,
            output_file=output_file if output_file else None,
            compile_flags=compile_flags if compile_flags else None,
            working_dir=working_dir
        )
        
        if result['success']:
            return f"""Compilation successful!
Source: {result['source_file']}
Compiler: {result['compiler']}
Output: {result['output_file']}
Command: {result['compile_command']}
Working directory: {result['working_dir']}

Standard output:
{result['stdout']}

Use execute_program to run the compiled code."""
        else:
            return f"""Compilation failed!
Source: {result['source_file']}
Compiler: {result['compiler']}
Command: {result['compile_command']}
Exit code: {result['exit_code']}

Standard error:
{result['stderr']}

Standard output:
{result['stdout']}"""
            
    except Exception as e:
        logger.error(f"Code compilation failed: {str(e)}")
        return f"Code compilation failed: {str(e)}"


@mcp.tool
def execute_program(connection_id: str, program_path: str, 
                   arguments: str = "", working_dir: str = "~") -> str:
    """Execute a compiled program or Python script
    
    Args:
        connection_id: SSH connection identifier
        program_path: Path to executable or Python script
        arguments: Command line arguments (optional)
        working_dir: Working directory for execution (default: ~)
        
    Returns:
        Program execution result and output
    """
    try:
        if not compiler_manager:
            raise Exception("Compiler manager not initialized")
        
        result = compiler_manager.execute_program(
            connection_id=connection_id,
            program_path=program_path,
            arguments=arguments if arguments else None,
            working_dir=working_dir
        )
        
        return f"""Program execution completed!
Program: {result['program_path']}
Arguments: {result['arguments']}
Command: {result['execution_command']}
Exit code: {result['exit_code']}
Working directory: {result['working_dir']}

Standard output:
{result['stdout']}

Standard error:
{result['stderr']}

Execution {'successful' if result['success'] else 'failed'}."""
        
    except Exception as e:
        logger.error(f"Program execution failed: {str(e)}")
        return f"Program execution failed: {str(e)}"


@mcp.tool
def submit_slurm_job(connection_id: str, script_content: str, job_name: str,
                    partition: str = "default", nodes: int = 1, cpus_per_task: int = 1,
                    memory: str = "4G", time_limit: str = "01:00:00") -> str:
    """Submit a Slurm batch job
    
    Args:
        connection_id: SSH connection identifier
        script_content: Job script content to execute
        job_name: Name for the job
        partition: Slurm partition to use (default: default)
        nodes: Number of nodes requested (default: 1)
        cpus_per_task: CPUs per task (default: 1)
        memory: Memory allocation, e.g., "4G" (default: 4G)
        time_limit: Time limit in HH:MM:SS format (default: 01:00:00)
        
    Returns:
        Job submission result with job ID
    """
    try:
        if not slurm_manager:
            raise Exception("Slurm manager not initialized")
        
        job_id = slurm_manager.submit_job(
            connection_id=connection_id,
            script_content=script_content,
            job_name=job_name,
            partition=partition,
            nodes=nodes,
            cpus_per_task=cpus_per_task,
            memory=memory,
            time_limit=time_limit
        )
        
        return f"""Slurm job submitted successfully!
Job ID: {job_id}
Job Name: {job_name}
Partition: {partition}
Nodes: {nodes}
CPUs per task: {cpus_per_task}
Memory: {memory}
Time limit: {time_limit}

Use check_slurm_job_status with job ID {job_id} to monitor the job."""
        
    except Exception as e:
        logger.error(f"Slurm job submission failed: {str(e)}")
        return f"Slurm job submission failed: {str(e)}"


@mcp.tool
def check_slurm_job_status(connection_id: str, job_id: str) -> str:
    """Check status of a specific Slurm job
    
    Args:
        connection_id: SSH connection identifier
        job_id: Slurm job ID to check
        
    Returns:
        Current job status and information
    """
    try:
        if not slurm_manager:
            raise Exception("Slurm manager not initialized")
        
        status_info = slurm_manager.check_job_status(connection_id, job_id)
        
        return f"""Slurm Job Status:
Job ID: {status_info.get('job_id', job_id)}
Status: {status_info.get('status', 'UNKNOWN')}
Job Name: {status_info.get('job_name', 'N/A')}
Partition: {status_info.get('partition', 'N/A')}
User: {status_info.get('user', 'N/A')}
Nodes: {status_info.get('nodes', 'N/A')}
Time Used: {status_info.get('time_used', 'N/A')}
Reason: {status_info.get('reason', 'N/A')}
Exit Code: {status_info.get('exit_code', 'N/A')}
Elapsed Time: {status_info.get('elapsed_time', 'N/A')}

Use wait_for_job_completion to wait for job completion."""
        
    except Exception as e:
        logger.error(f"Job status check failed: {str(e)}")
        return f"Job status check failed: {str(e)}"


@mcp.tool
def execute_sequential_commands(connection_id: str = "", commands: str = "", 
                              working_dir: str = "~", 
                              stop_on_error: bool = True) -> str:
    """Execute multiple commands sequentially in the same session
    
    Args:
        connection_id: SSH connection identifier (empty = auto-connect using .env)
        commands: Multi-line string with commands separated by newlines or semicolons
        working_dir: Initial working directory (default: "~")
        stop_on_error: Stop execution if any command fails (default: True)
        
    Returns:
        Combined output from all commands with execution status
        
    Example:
        commands = '''
        cd /path/to/project1
        make clean
        make build
        sbatch run.slurm
        '''
    """
    try:
        if not ssh_manager:
            raise Exception("SSH manager not initialized")
        
        # Get or create connection
        actual_connection_id = get_or_create_connection(connection_id)
        
        # Clean and split commands
        cmd_list = []
        for line in commands.strip().split('\n'):
            line = line.strip()
            if line and not line.startswith('#'):
                # Split on semicolons too
                for cmd in line.split(';'):
                    cmd = cmd.strip()
                    if cmd:
                        cmd_list.append(cmd)
        
        if not cmd_list:
            return "No valid commands provided"
        
        results = []
        current_dir = working_dir
        
        for i, command in enumerate(cmd_list, 1):
            try:
                # Handle cd commands specially to track directory changes
                if command.startswith('cd '):
                    new_dir = command[3:].strip()
                    if new_dir.startswith('/'):
                        current_dir = new_dir
                    elif new_dir == '~':
                        current_dir = '~'
                    elif new_dir == '..':
                        if current_dir != '~' and current_dir != '/':
                            current_dir = os.path.dirname(current_dir)
                    else:
                        if current_dir == '~':
                            current_dir = new_dir
                        else:
                            current_dir = os.path.join(current_dir, new_dir)
                
                # Execute the command
                exit_code, stdout, stderr = ssh_manager.execute_command(
                    actual_connection_id, command, current_dir
                )
                
                status = "SUCCESS" if exit_code == 0 else "FAILED"
                result = f"""Command {i}: {command}
Directory: {current_dir}
Status: {status} (exit code: {exit_code})
Output: {stdout.strip() if stdout else '(no output)'}
{f'Error: {stderr.strip()}' if stderr else ''}
{'='*50}"""
                
                results.append(result)
                
                # Stop on error if requested
                if stop_on_error and exit_code != 0:
                    results.append(f"Stopping execution due to error in command {i}")
                    break
                    
            except Exception as e:
                error_result = f"""Command {i}: {command}
Directory: {current_dir}
Status: EXCEPTION
Error: {str(e)}
{'='*50}"""
                results.append(error_result)
                
                if stop_on_error:
                    results.append(f"Stopping execution due to exception in command {i}")
                    break
        
        return f"""Sequential command execution completed:
Total commands: {len(cmd_list)}
Commands executed: {len([r for r in results if not r.startswith('Stopping')])}

{chr(10).join(results)}"""
        
    except Exception as e:
        logger.error(f"Sequential command execution failed: {str(e)}")
        return f"Sequential command execution failed: {str(e)}"


@mcp.tool
def wait_for_job_completion(connection_id: str, job_id: str, 
                           check_interval: int = 30, timeout: int = 3600) -> str:
    """Wait for a Slurm job to complete with timeout
    
    Args:
        connection_id: SSH connection identifier
        job_id: Slurm job ID to wait for
        check_interval: Time between status checks in seconds (default: 30)
        timeout: Maximum wait time in seconds (default: 3600)
        
    Returns:
        Final job status and output when completed
    """
    try:
        if not slurm_manager:
            raise Exception("Slurm manager not initialized")
        
        # Run async function in sync context
        loop = asyncio.new_event_loop()
        asyncio.set_event_loop(loop)
        try:
            final_status = loop.run_until_complete(
                slurm_manager.wait_for_job_completion(
                    connection_id, job_id, check_interval, timeout
                )
            )
        finally:
            loop.close()
        
        return f"""Job completion result:
Job ID: {final_status.get('job_id', job_id)}
Final Status: {final_status.get('status', 'UNKNOWN')}
Job Name: {final_status.get('job_name', 'N/A')}
Exit Code: {final_status.get('exit_code', 'N/A')}
Elapsed Time: {final_status.get('elapsed_time', 'N/A')}

Job Output (last 20 lines):
{final_status.get('stdout', 'No output available')}

Job Errors (last 20 lines):
{final_status.get('stderr', 'No errors')}

Job completed {'successfully' if final_status.get('status') == 'COMPLETED' else 'with issues'}."""
        
    except Exception as e:
        logger.error(f"Job wait failed: {str(e)}")
        return f"Job wait failed: {str(e)}"


@mcp.tool
def upload_file(connection_id: str = "", local_path: str = "", remote_path: str = "") -> str:
    """Upload file from local to remote server
    
    Args:
        connection_id: SSH connection identifier (empty = auto-connect using .env)
        local_path: Path to local file to upload
        remote_path: Destination path on remote server
        
    Returns:
        Upload status message
    """
    try:
        if not ssh_manager:
            raise Exception("SSH manager not initialized")
        
        # Get or create connection
        actual_connection_id = get_or_create_connection(connection_id)
        
        success = ssh_manager.upload_file(actual_connection_id, local_path, remote_path)
        
        if success:
            return f"Successfully uploaded {local_path} to {actual_connection_id}:{remote_path}"
        else:
            return f"Failed to upload {local_path} to {actual_connection_id}:{remote_path}"
            
    except Exception as e:
        logger.error(f"File upload failed: {str(e)}")
        return f"File upload failed: {str(e)}"


@mcp.tool
def download_file(connection_id: str = "", remote_path: str = "", local_path: str = "") -> str:
    """Download file from remote server to local
    
    Args:
        connection_id: SSH connection identifier (empty = auto-connect using .env)
        remote_path: Path to file on remote server
        local_path: Local destination path
        
    Returns:
        Download status message
    """
    try:
        if not ssh_manager:
            raise Exception("SSH manager not initialized")
        
        # Get or create connection
        actual_connection_id = get_or_create_connection(connection_id)
        
        success = ssh_manager.download_file(actual_connection_id, remote_path, local_path)
        
        if success:
            return f"Successfully downloaded {actual_connection_id}:{remote_path} to {local_path}"
        else:
            return f"Failed to download {actual_connection_id}:{remote_path} to {local_path}"
            
    except Exception as e:
        logger.error(f"File download failed: {str(e)}")
        return f"File download failed: {str(e)}"


@mcp.tool
def list_remote_files(connection_id: str = "", directory: str = "~", pattern: str = "*") -> str:
    """List files in a remote directory
    
    Args:
        connection_id: SSH connection identifier (empty = auto-connect using .env)
        directory: Remote directory path (default: ~)
        pattern: File pattern filter (default: *)
        
    Returns:
        List of files and directories with details
    """
    try:
        if not ssh_manager:
            raise Exception("SSH manager not initialized")
        
        # Get or create connection
        actual_connection_id = get_or_create_connection(connection_id)
        
        files = ssh_manager.list_remote_files(actual_connection_id, directory, pattern)
        
        if not files:
            return f"No files found in {actual_connection_id}:{directory} matching pattern '{pattern}'"
        
        result = f"Files in {actual_connection_id}:{directory} (pattern: {pattern}):\n\n"
        for file_info in files:
            result += f"{file_info['permissions']} {file_info['owner']} {file_info['group']} "
            result += f"{file_info['size']:>8} {file_info['date']} {file_info['name']}\n"
        
        result += f"\nTotal: {len(files)} items"
        return result
        
    except Exception as e:
        logger.error(f"File listing failed: {str(e)}")
        return f"File listing failed: {str(e)}"


# Additional helper tools
@mcp.tool
def list_connections() -> str:
    """List all active SSH connections
    
    Returns:
        List of active connections with details
    """
    try:
        if not ssh_manager:
            return "SSH manager not initialized"
        
        connections = ssh_manager.list_connections()
        
        if not connections:
            return "No active SSH connections"
        
        result = "Active SSH Connections:\n\n"
        for conn in connections:
            result += f"Connection ID: {conn['connection_id']}\n"
            result += f"  Host: {conn['host']}\n"
            result += f"  Username: {conn['username']}\n"
            result += f"  Port: {conn['port']}\n"
            result += f"  Connected at: {conn['connected_at']}\n"
            result += f"  Status: {'Active' if conn['is_active'] else 'Inactive'}\n\n"
        
        return result
        
    except Exception as e:
        logger.error(f"Connection listing failed: {str(e)}")
        return f"Connection listing failed: {str(e)}"


@mcp.tool
def get_compiler_info(connection_id: str = "", compiler: str = "") -> str:
    """Get compiler version and availability information
    
    Args:
        connection_id: SSH connection identifier (empty = auto-connect using .env)
        compiler: Specific compiler to check (empty = check all)
        
    Returns:
        Compiler information and availability
    """
    try:
        if not compiler_manager:
            raise Exception("Compiler manager not initialized")
        
        # Get or create connection
        actual_connection_id = get_or_create_connection(connection_id)
        
        if compiler:
            info = compiler_manager.get_compiler_info(actual_connection_id, compiler)
            if info['available']:
                return f"Compiler {compiler} is available on {actual_connection_id}:\n{info['version_output']}"
            else:
                return f"Compiler {compiler} is not available on {actual_connection_id}: {info.get('error', 'Unknown error')}"
        else:
            compilers = compiler_manager.list_available_compilers(actual_connection_id)
            result = f"Compiler availability on {actual_connection_id}:\n\n"
            for comp in compilers:
                status = "✓ Available" if comp['available'] else "✗ Not available"
                result += f"{comp['compiler']}: {status}\n"
                if comp['available']:
                    result += f"  Version: {comp['version_output'].strip()}\n"
                else:
                    result += f"  Error: {comp.get('error', 'Unknown')}\n"
                result += "\n"
            
            return result
            
    except Exception as e:
        logger.error(f"Compiler info check failed: {str(e)}")
        return f"Compiler info check failed: {str(e)}"


def cleanup_on_exit():
    """Cleanup function to run on server shutdown"""
    if ssh_manager:
        ssh_manager.cleanup()
    logger.info("SSH Slurm MCP Server cleanup completed")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="SSH Slurm MCP Server")
    parser.add_argument("--transport", default="stdio", help="Transport method")
    args = parser.parse_args()
    
    try:
        # Initialize managers
        initialize_managers()
        
        # Register cleanup function
        import atexit
        atexit.register(cleanup_on_exit)
        
        logger.info("Starting SSH Slurm MCP Server...")
        
        # Run MCP server
        mcp.run(transport=args.transport)
        
    except KeyboardInterrupt:
        logger.info("Server interrupted by user")
        cleanup_on_exit()
    except Exception as e:
        logger.error(f"Server error: {str(e)}")
        cleanup_on_exit()
        sys.exit(1)
